package com.dingzhen.support;

import com.dingzhen.convert.JacksonJsonMessageConvert;
import com.dingzhen.convert.MessageConvert;
import com.dingzhen.entity.Page;
import com.dingzhen.entity.SearchScrollHits;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.search.ClearScrollRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.index.reindex.UpdateByQueryRequest;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 *  Template Methods for operation Elasticsearch
 */
public class ElasticsearchRestTemplate {
        private final RestHighLevelClient restHighLevelClient;
        private final MessageConvert messageConvert;

        public ElasticsearchRestTemplate(RestHighLevelClient restHighLevelClient) {
            this.restHighLevelClient = restHighLevelClient;
            this.messageConvert = new JacksonJsonMessageConvert();
        }

        public <T> Page<T> queryForPage(SearchSourceBuilder searchBuilder, String indexName, Class<T> clazz, int pageNum, int pageSize) {
            searchBuilder.from((pageNum - 1) * pageSize).size(pageSize);
            SearchResponse response = this.doQuery(searchBuilder, indexName);
            List<T> list = messageConvert.convert(response.getHits(), clazz);
            return new Page<>(list, pageNum, pageSize, response.getHits().getTotalHits().value);
        }

        public <T> Page<T> queryForPage(SearchSourceBuilder searchBuilder, String indexName, Class<T> clazz) {
            return this.queryForPage(searchBuilder, indexName, clazz, 1, 10);
        }

        public <T> T queryForObject(SearchSourceBuilder searchBuilder, String indexName, Class<T> clazz) {
            Page<T> page = queryForPage(searchBuilder, indexName, clazz);
            return page.getTotalElements() > 0 ? page.getContent().get(0) : null;
        }

        public <T> List<T> queryForList(SearchSourceBuilder searchBuilder, String indexName, Class<T> clazz) {
            SearchResponse response = this.doQuery(searchBuilder, indexName);
            return messageConvert.convert(response.getHits(), clazz);
        }

        public long count(SearchSourceBuilder searchBuilder, String indexName) {

            CountRequest countRequest = new CountRequest();
            CountResponse count;
            try {
                count = restHighLevelClient.count(countRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
            return count.getCount();
        }

        public SearchResponse searchScroll(SearchSourceBuilder searchBuilder, long scrollTimeInMillis, String indexName) {

            SearchRequest searchRequest = new SearchRequest(indexName);
            searchRequest.scroll(TimeValue.timeValueMillis(scrollTimeInMillis));
            searchRequest.source(searchBuilder);
            try {
                return restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }

        public <T> SearchScrollHits<T> searchScrollStart(SearchSourceBuilder searchSourceBuilder, long scrollTimeInMillis, String indexName, Class<T> clazz) {
            SearchResponse response = searchScroll(searchSourceBuilder, scrollTimeInMillis, indexName);
            return wrapResponse(response, clazz);
        }

        private SearchResponse searchScrollContinue(String scrollId) {
            SearchScrollRequest searchScrollRequest = new SearchScrollRequest(scrollId);
            SearchResponse response;
            try {
                response = restHighLevelClient.scroll(searchScrollRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
            return response;
        }

        public <T> SearchScrollHits<T> searchScrollContinue(String scrollId, Class<T> clazz) {
            SearchResponse response = searchScrollContinue(scrollId);
            return wrapResponse(response, clazz);
        }

        private <T> SearchScrollHits<T> wrapResponse(SearchResponse response, Class<T> clazz) {
            SearchScrollHits<T> searchScrollHits = new SearchScrollHits<>(response);
            searchScrollHits.setHits(messageConvert.convert(response.getHits(), clazz));
            return searchScrollHits;
        }

        public void searchScrollClear(String... scrollIds) {
            ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
            clearScrollRequest.scrollIds(Arrays.asList(scrollIds));
            try {
                restHighLevelClient.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }

        private SearchResponse doQuery(SearchSourceBuilder searchSourceBuilder, String indexName) {
            SearchRequest searchRequest = new SearchRequest(indexName);
            searchRequest.source(searchSourceBuilder);
            try {
                return restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }

        public void createIndex(String indexName) {
            CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName);
            try {
                restHighLevelClient.indices().create(createIndexRequest,RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }

        }

        public void createIndex(String indexName, Map<String,Object> settings){
            CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName);
            createIndexRequest.settings(settings);
            try {
                restHighLevelClient.indices().create(createIndexRequest,RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }
        public void createIndex(String indexName, Map<String,Object> settings,Map<String,Object> mapping){
            CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName);
            createIndexRequest.settings(settings);
            createIndexRequest.mapping(mapping);
            try {
                restHighLevelClient.indices().create(createIndexRequest,RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }

        public void createIndex(String indexName,String mappings) {
            CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName);
            createIndexRequest.mapping(mappings, XContentType.JSON);
            try {
                restHighLevelClient.indices().create(createIndexRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }

        }
        public void createIndexIfNotCreated(String indexName) {
            if (!indexExists(indexName)) {
                createIndex(indexName);
            }
        }
        public void deleteIndex(String... indices){
            DeleteIndexRequest deleteIndexRequest = new DeleteIndexRequest(indices);
            try {
                restHighLevelClient.indices().delete(deleteIndexRequest,RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }

        public boolean indexExists(String indexName) {
            GetIndexRequest getIndexRequest = new GetIndexRequest(indexName);
            try {
                return restHighLevelClient.indices().exists(getIndexRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }
        public void saveOrUpdate(String id, String docAsJson, String indexName) {
            UpdateRequest updateRequest = new UpdateRequest(indexName,id);
            updateRequest.doc(docAsJson,XContentType.JSON);
            updateRequest.docAsUpsert(true);
            try {
                restHighLevelClient.update(updateRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }

        public void deleteDocById(String id, String indexName) {
            //
            DeleteRequest deleteRequest = new DeleteRequest(indexName, id);
            try {
                restHighLevelClient.delete(deleteRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }


        public void deleteByQuery(QueryBuilder queryBuilder, String indexName) {
            //
            DeleteByQueryRequest deleteRequest = new DeleteByQueryRequest(indexName);
            deleteRequest.setQuery(queryBuilder);
            try {
                restHighLevelClient.deleteByQuery(deleteRequest, RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }
        }

        public void updateByQuery(QueryBuilder queryBuilder, ScriptBuilder builder,String indexName){
            UpdateByQueryRequest updateByQueryRequest = new UpdateByQueryRequest(indexName);
            updateByQueryRequest.setQuery(queryBuilder);
            updateByQueryRequest.setScript(builder.build());
            try {
                restHighLevelClient.updateByQuery(updateByQueryRequest,RequestOptions.DEFAULT);
            } catch (IOException e) {
                throw new ElasticsearchException(e);
            }

        }
    }
